"""
/* Copyright 2018 The Enflame Tech Company. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
"""
# !/usr/bin/python
# coding=utf-8

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import random
import sys
import numpy as np
import signal
import os
from collections import OrderedDict
from datetime import datetime
from environs import Env
from tensorflow.python.client import timeline
from tensorflow.core.framework import summary_pb2
import os

FLAGS = tf.app.flags.FLAGS
env = Env()

from config.dtypes import dtype_mapping
from config.sizes import size_mapping
from config.datasets import dataset_mapping
from utils.dtu_utils import *
from utils.host_info import HostInfo
from utils.sw_info import SWInfo
from utils.hw_info import HWInfo
from utils.dtu_logger import LOGGER as logger
from utils.flags import InfoDict
from utils.time_watcher import EnflameThread
import classifier_main

ROOT_PATH = os.path.dirname(os.path.abspath(__file__))
LOG_DIR = "{}/logs".format(ROOT_PATH)
if not os.path.exists(LOG_DIR):
    try:
        os.makedirs(LOG_DIR)
    except OSError as ex:
        logger.error("Error message: {}".format(ex))
REPORT_DIR = "{}/test_report".format(ROOT_PATH)
if not os.path.exists(REPORT_DIR):
    try:
        os.makedirs(REPORT_DIR)
    except OSError as ex:
        logger.error("Error message: {}".format(ex))
CKPT_DIR = "./models/checkpoints"
if not os.path.exists(CKPT_DIR):
    try:
        os.makedirs(CKPT_DIR)
    except OSError as ex:
        logger.error("Error message: {}".format(ex))
EVENT_DIR = "{}/models/tensorboard".format(ROOT_PATH)
if not os.path.exists(EVENT_DIR):
    try:
        os.makedirs(EVENT_DIR)
    except OSError as ex:
        logger.error("Error message: {}".format(ex))

# globa constant
NUM_CHANNEL = 3
NUM_BETWEEN_EVAL = FLAGS.epochs_between_evals


def quit_gracefully(*args):
    logger.error('caught signal ,quit')
    exit(0)


def main(_):

    signal.signal(signal.SIGINT, quit_gracefully)

    # collect information
    try:
        info_dict = InfoDict()
        base_info = info_dict.base_info
        opt_info = info_dict.opt_info
        config_info = info_dict.config_info
        test_info = info_dict.test_info
        func_info = info_dict.func_info
        json_report(func_info)
        host_info = HostInfo().dict
        hw_info = HWInfo(test_info['device']).dict
        sw_info = SWInfo().dict
        test_info.update(host_info)
        test_info.update(sw_info)
        test_info.update(hw_info)
        json_report(test_info)
        test_info.update(base_info)
    except Exception as ex:
        logger.warn("Collect information failed")

    global_watcher = EnflameThread().watcher(batch_size_per_step=test_info['batch_size'] * test_info['num_cluster'])

    global_watcher.watch("import horovod start")
    try:
        """Horovod will be used once successfully imported and initialized."""
        import horovod.tensorflow as hvd
        hvd.init()
        is_chief = hvd.local_rank() == 0
        hvd_info = OrderedDict([('enable_horovod', env.bool("ENABLE_HOROVOD", True)),
                                ('hvd_size', hvd.size()),
                                ('local_rank', hvd.local_rank()),
                                ('rank', hvd.rank())
                                ])
        test_info.update(hvd_info)
    except ImportError as e:
        logger.warn("Horovod init failed! Error message: {}".format(e))
    global_watcher.watch("import horovod end")

    # set debug seed
    if FLAGS.debug_mode:
        seed = 123
        random.seed(seed)
        tf.set_random_seed(seed)
        np.random.seed(seed)
        np.set_printoptions(threshold=int(sys.maxsize))

    if FLAGS.data_dir:
        dataset_location = FLAGS.data_dir
    else:
        dataset_location = "{}/dataset/{}".format(ROOT_PATH, test_info['dataset'])

    # global check
    global_watcher.watch("global check start")
    assert test_info['device'] in ['cpu', 'gpu', 'dtu', 'tpu', 'xla_gpu', 'xla_cpu'], 'Unsupported hardware: {}'.format(
        FLAGS.device.lower())

    try:
        dtype = dtype_mapping[test_info['dtype']]
    except Exception as exception:
        logger.error("Unsupported data dtype: {}".format(test_info['dtype']))
        sys.exit(6)

    try:
        output_size = size_mapping[test_info['model']]
        test_info['output_size'] = output_size
        test_info['num_channels'] = NUM_CHANNEL
    except Exception as exception:
        logger.error("Unsupported model name: {}".format(test_info['model']))
        sys.exit(6)

    try:
        num_class = dataset_mapping[test_info['dataset']]['class_number']
        test_info['num_class'] = num_class
    except Exception as exception:
        logger.error("Unsupported dataset: {}".format(test_info['dataset']))
        sys.exit(6)
    global_watcher.watch("global check end")

    # json_report(test_info)

    def create_config():
        config = tf.ConfigProto(allow_soft_placement=config_info['allow_soft_placement'],
                                log_device_placement=config_info['log_device_placement'],
                                inter_op_parallelism_threads=config_info['inter_op_parallelism_threads'],
                                intra_op_parallelism_threads=config_info['intra_op_parallelism_threads'])
        if test_info['device'] in ['gpu', 'xla_gpu']:
            if FLAGS.xla_jit:
                config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
            else:
                config.gpu_options.allow_growth = True
                config.gpu_options.per_process_gpu_memory_fraction = 0.90
            if test_info['enable_horovod']:
                config.gpu_options.visible_device_list = str(hvd.local_rank())
        elif test_info['device'] in ['dtu', 'xla_dtu']:
            config.dtu_options.visible_device_list = str(test_info['local_rank'])
            logger.info("config.dtu_options.visible_device_list={}".format(config.dtu_options.visible_device_list))
        else:
            logger.info("default config options")

        from tensorflow.core.protobuf import rewriter_config_pb2
        off = rewriter_config_pb2.RewriterConfig.OFF
        config.graph_options.rewrite_options.memory_optimization = off
        return config

    if FLAGS.enable_profiler:
        options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
        run_metadata = tf.RunMetadata()
        step_run_metadata = tf.RunMetadata()
    else:
        options = None
        step_run_metadata = None

    bench = classifier_main.Benchmark(test_info=test_info, dtype=dtype)
    global_step = tf.train.get_or_create_global_step()

    # data processing
    global_watcher.watch("data processing start")
    with tf.device(bench.cpu_device):
        if FLAGS.is_training:
            train_data_dir = "{}/train".format(dataset_location)
            iterator_train = bench._generate_data(data_dir=train_data_dir, is_training=True)
            iterator_initial_op = iterator_train.initializer
            if FLAGS.enable_evaluate:
                eval_data_dir = "{}/evaluate".format(dataset_location)
                iterator_eval = bench._generate_data(data_dir=eval_data_dir, is_training=False)
                iterator_initial_op = tf.group(iterator_eval.initializer, iterator_initial_op)

        else:
            infer_data_dir = '{}/test'.format(dataset_location)
            iterator_infer = bench._generate_data(data_dir=infer_data_dir, is_training=False)
            iterator_initial_op = iterator_infer.initializer
    global_watcher.watch("data processing end")

    # build graph
    global_watcher.watch("build graph start")
    if FLAGS.xla_jit:
        with tf.contrib.compiler.jit.experimental_jit_scope() as jit_scope:
            if FLAGS.is_training:
                train_op, accuracy_op, loss_op = bench.build_training_graph(iterator=iterator_train,
                                                                            dropout_rate=FLAGS.dropout_rate,
                                                                            global_step=global_step)
                if FLAGS.enable_evaluate:
                    eval_accuracy_op, eval_accuracy_top_5_op = bench.build_evaluate_graph(iterator=iterator_eval,
                                                                                          dropout_rate=1.0)
            else:
                predict_op = bench.build_inference_graph(iterator=iterator_infer)
    else:
        with tf.device(bench.default_device):
            if FLAGS.is_training:
                train_op, accuracy_op, loss_op = bench.build_training_graph(iterator=iterator_train,
                                                                            dropout_rate=FLAGS.dropout_rate,
                                                                            global_step=global_step)
                tf.summary.scalar('loss', loss_op)
                if FLAGS.enable_evaluate:
                    eval_accuracy_op, eval_accuracy_top_5_op = bench.build_evaluate_graph(iterator=iterator_eval,
                                                                                          dropout_rate=1.0)

            else:
                predict_op = bench.build_inference_graph(iterator=iterator_infer)
    global_watcher.watch("build graph end")
    if FLAGS.enable_dump_graph:
        graph = tf.get_default_graph()
        graph_path = "{}/{}_on_{}_before_session.json".format(LOG_DIR, test_info['model'], test_info['device'])
        export_graph_to_file(graph, graph_path)
        tf.io.write_graph(graph, LOG_DIR, '{}_on_{}_before_session_{}.pbtxt'.format(test_info['model'], test_info['device'],
                                                                                    datetime.now().strftime('%Y%m%d%H%M%S.%f')[:-3]))

    merged_summary = tf.summary.merge_all()

    # session process
    with tf.device(bench.default_device):
        with tf.Session(config=create_config()) as sess:
            # initial
            global_watcher.watch("variable init start")
            sess.run(bench.variables_initial_op())
            sess.run(iterator_initial_op)
            global_watcher.watch("variable init end")

            if FLAGS.enable_dump_graph:
                graph = sess.graph
                graph_path = "{}/{}_on_{}_after_init.json".format(LOG_DIR, test_info['model'], test_info['device'])
                tf.io.write_graph(sess.graph.as_graph_def(add_shapes=True), LOG_DIR, '{}_on_{}_after_init_{}.pbtxt'.format(test_info['model'],test_info['device'],
                                                                                                 datetime.now().strftime('%Y%m%d%H%M%S.%f')[:-3]))
                export_graph_to_file(graph, graph_path)
                tf.io.write_graph(sess.graph_def, '/tmp/', 'train_{}_{}_after_init.pbtxt'.format(test_info['model'],
                                                                                      datetime.now().strftime('%Y%m%d%H%M%S.%f')[:-3]))

            # create summary writer
            writer = tf.summary.FileWriter(logdir=EVENT_DIR + '/train-{}'.format(bench.params['local_rank']), graph=sess.graph, filename_suffix='.train')
            writer.add_graph(sess.graph)
            if FLAGS.enable_evaluate:
                eval_writer = tf.summary.FileWriter(logdir=EVENT_DIR + '/eval-{}'.format(bench.params['local_rank']), graph=sess.graph,
                                                    filename_suffix='.eval')

            if FLAGS.enable_saver or not FLAGS.is_training:
                saver = tf.train.Saver(max_to_keep=100)
                ckpt = tf.train.latest_checkpoint(CKPT_DIR,
                                                  latest_filename="checkpoint_{}_{}".format(FLAGS.model, FLAGS.dataset))
                if ckpt is not None:
                    saver.restore(sess, ckpt)

            # horovod broadcast
            if bench.params['enable_horovod']:
                global_watcher.watch("horovod broadcast start")
                sess.run(hvd.broadcast_global_variables(0))
                global_watcher.watch("horovod broadcast end")

            for epoch in range(1, test_info['epoch']+1):
                if FLAGS.is_training:
                    for local_step in range(bench._training_step_num_per_epoch()):
                        global_watcher.watch("epoch:{}, local_step:{}, session_start".format(epoch, local_step))
                        _, loss, step_count = sess.run([train_op, loss_op, global_step], options=options, run_metadata=step_run_metadata)
                        global_watcher.watch("epoch:{}, local_step:{}, session_end".format(epoch, local_step))
                        if (local_step + 1) % FLAGS.display_step == 0:
                            accuracy, s= sess.run([accuracy_op, merged_summary])
                            logger.info("loss: {}; accuracy: {}".format(loss, accuracy))
                            writer.add_summary(s, step_count)
                            assert not check_valid_data(loss), 'Found valid data in loss, NAN or INF'
                            assert not check_valid_data(accuracy[1]), 'Found valid data in accuracy, NAN or INF'
                        if FLAGS.enable_profiler:
                            run_metadata.MergeFrom(step_run_metadata)
                        # 新加代码
                        if local_step == int(os.environ['LOG_AFTER_STEP_NUM']):
                            # script_name, device, dtype, model, optimizer, batch_size, fps
                            info = 'enflame_model_test,{},{},{},{},{},{}'.format(
                                FLAGS.device,
                                FLAGS.dtype,
                                FLAGS.model,
                                FLAGS.optimizer,
                                global_watcher.batch_size_per_step, # 如果是dtu，则有多个核
                                sum(global_watcher.step_fps_list) / len(global_watcher.step_fps_list),
                            )
                            with open(os.environ['LOG_FILE'], 'a+') as f:
                                f.write(info)
                                f.write('\n')
                            sys.exit(0)
                        # 新加代码
                    if FLAGS.enable_evaluate and (epoch % NUM_BETWEEN_EVAL == 0 or epoch == 1):
                        tower_eval_accuracy = 0
                        tower_eval_accuracy_top_5 = 0
                        for local_step in range(bench._evaluate_step_num_per_epoch()):
                            eval_accuracy, eval_accuracy_top_5 = sess.run([eval_accuracy_op, eval_accuracy_top_5_op])
                            logger.info("eval_accuracy: {}; eval_accuracy_top_5: {}".format(eval_accuracy, eval_accuracy_top_5))
                            tower_eval_accuracy += eval_accuracy
                            tower_eval_accuracy_top_5 += eval_accuracy_top_5
                        avg_eval_accuracy = tower_eval_accuracy / bench._evaluate_step_num_per_epoch()
                        avg_eval_accuracy_top_5 = tower_eval_accuracy_top_5 / bench._evaluate_step_num_per_epoch()
                        logger.info("avg_eval_accuracy: {}; avg_eval_accuracy_top_5: {}".format(avg_eval_accuracy, avg_eval_accuracy_top_5))
                        s = summary_pb2.Summary(
                            value=[summary_pb2.Summary.Value(tag="eval_accuracy", simple_value=avg_eval_accuracy)])
                        s_5 = summary_pb2.Summary(
                             value=[summary_pb2.Summary.Value(tag="eval_accuracy_top_5", simple_value=avg_eval_accuracy_top_5)])
                        eval_writer.add_summary(s, step_count)
                        eval_writer.add_summary(s_5, step_count)

                    if epoch % NUM_BETWEEN_EVAL == 0 and FLAGS.enable_saver and bench.params['local_rank'] == 0:
                        saver.save(sess=sess, save_path="{}/enflame_{}_{}".format(CKPT_DIR, FLAGS.model, FLAGS.dataset),
                                   global_step=global_step,
                                   latest_filename="checkpoint_{}_{}".format(FLAGS.model, FLAGS.dataset))

                else:
                    global_watcher.watch("epoch:{}, session_start".format(epoch))
                    predict, s, step_count = sess.run([predict_op, merged_summary, global_step], options=options,
                                                      run_metadata=step_run_metadata)
                    global_watcher.watch("epoch:{}, session end".format(epoch))
                    logger.info(
                        "predict class: {}, real class: {}; accuracy: {}".format(predict['predict class'],
                                                                                 predict['real class'],
                                                                                 calc_inference_accuracy(
                                                                                     predict['predict class'],
                                                                                     predict['real class'])))
                    # logger.info("probability: {}".format(predict['prob']))
                    if FLAGS.enable_profiler:
                        run_metadata.MergeFrom(step_run_metadata)

                    if bench.params['local_rank'] == 0:
                        writer.add_summary(s, step_count)

    with tf.device(bench.cpu_device):
        # profile
        if FLAGS.enable_profiler:
            global_watcher.watch("Profile start")
            profile_log = "{}/{}_{}_profile.log".format(REPORT_DIR, bench.params['model'], bench.params['device'])
            profile_opts = tf.profiler.ProfileOptionBuilder(
                tf.profiler.ProfileOptionBuilder.time_and_memory()).with_file_output(profile_log).build()
            tf.profiler.profile(tf.get_default_graph(), run_meta=run_metadata, cmd='op', options=profile_opts)
            profile_opts = tf.profiler.ProfileOptionBuilder.float_operation()
            tf.profiler.profile(tf.get_default_graph(), run_meta=run_metadata, cmd='op', options=profile_opts)
            chrome_trace = timeline.Timeline(run_metadata.step_stats).generate_chrome_trace_format()
            profile_json = "{}/{}_{}_timeline.json".format(REPORT_DIR, bench.params['model'],
                                                           bench.params['device'])
            with open(profile_json, 'w') as f:
                f.write(chrome_trace)
            global_watcher.watch("Profile end")
        # update all information
        runtime_info = OrderedDict([('training_step_per_epoch', bench._training_step_num_per_epoch()),
                                    ('evaluate_step_per_epoch', bench._evaluate_step_num_per_epoch()),
                                    ('fps', global_watcher.calc_fps()),
                                    ('min_fps', global_watcher.min_fps()),
                                    ('max_fps', global_watcher.max_fps()),
                                    ('session_duration', sum(global_watcher.duration_time_list)),
                                    ('total_duration', (datetime.now() - test_info['start_time']).total_seconds()),
                                    ('start_time', test_info['start_time'].strftime("%Y-%m-%d %H:%M:%S"))
                                    ])
        test_info.update(runtime_info)
        json_report(test_info)

if __name__ == '__main__':
    tf.app.run()
